import pandas as pd
from mlxtend.preprocessing import TransactionEncoder
from mlxtend.frequent_patterns import apriori,association_rules
import matplotlib.pyplot as plt
import networkx as nx

df = pd.read_excel('associationRules/餐厅数据.xlsx')

#print(df.head)
trsations = df['菜品'].str.split(',').to_list()

ts=TransactionEncoder()
te_ary=ts.fit(trsations).transform(trsations)
df_enecoded=pd.DataFrame(te_ary,columns=ts.columns_)

#print(df_enecoded.head)
frequent_itemsets=apriori(df_enecoded,min_support=0.1,use_colnames=True)
frequent_itemsets.sort_values(by='support',ascending=False,inplace=True)

#print(frequent_itemsets[frequent_itemsets.itemsets.apply(lambda x : len(x) == 2)])

rules=association_rules(frequent_itemsets,metric='confidence',min_threshold=0.1)

effective = rules[
    (rules['lift'] > 1) & (rules['confidence'] > 0.1)
].sort_values(by=['lift','confidence'],ascending=False)

frequent_itemsets.to_pickle("frequent_itemsets.pkl")
effective.to_pickle("rule.pkl")